from .ai_unit import AIUnit, AIUnitParams, TrainningParam
import websocket
from typing import Dict
import threading

unit_saver: Dict[websocket.WebSocketApp, AIUnit] = {}
rlock = threading.RLock()


def on_message(ws: websocket.WebSocketApp, message: str):
    unit = unit_saver.get(ws)
    if unit is not None:
        rlock.acquire()
        unit.process_connect_message(message)
        rlock.release()
    else:
        print("unit None")
        ws.close()


def on_error(ws, error):
    print(ws)
    print(error)


def on_close(ws):
    print(ws)
    print("### closed ###")


def on_open(ws):
    print("on_open: ", unit_saver.get(ws).url)


ai_unit_params = AIUnitParams(5)
trainning_params = TrainningParam(0.9, 0.6, 32)


class WebsocketConnector:
    url: str
    ws: websocket.WebSocketApp

    def __init__(self, url: str):
        self.url = url
        websocket.enableTrace(True)
        self.ws = websocket.WebSocketApp(url,
                                         on_open=on_open,
                                         on_message=on_message,
                                         on_error=on_error,
                                         on_close=on_close)
        unit_saver[self.ws] = AIUnit(self.url, self.ws, ai_unit_params, trainning_params)
        self.ws.run_forever()

    def close(self):
        print("close")
        self.ws.close()
